import numpy as np


class Optimizer:
    def __init__(self, params):
        self.params = [x for x in params if x._trainable]

    def zero_grad(self):
        for param in self.params:
            param.grad = np.zeros_like(param.data)


class SGD(Optimizer):
    def __init__(self, params, lr=0.001):
        super().__init__(params)
        self.lr = lr

    def step(self):
        for t in self.params:
            t.data -= t.grad * self.lr


class Adam(Optimizer):
    def __init__(self, params, lr=0.001, b1=0.9, b2=0.999, eps=1e-8):
        super().__init__(params)
        self.lr, self.b1, self.b2, self.eps, self.t = lr, b1, b2, eps, 0

        self.m = [np.zeros(t.shape) for t in self.params]
        self.v = [np.zeros(t.shape) for t in self.params]

    def step(self):
        self.t = self.t + 1
        a = self.lr * ((1.0 - self.b2**self.t)**0.5) / (1.0 - self.b1**self.t)
        for i, t in enumerate(self.params):
            self.m[i] = self.b1 * self.m[i] + (1.0 - self.b1) * t.grad
            self.v[i] = self.b2 * self.v[i] + (1.0 - self.b2) * t.grad * t.grad
            t -= a * self.m[i]/(self.v[i]**0.5 + self.eps)
